import _init_paths

import glob
import numpy as np
import argparse, os

from base64 import b64encode
from elasticsearch import Elasticsearch
from elasticsearch import helpers



import logging
logging.basicConfig(
    format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)
logger = logging.getLogger(__name__)


# MODEL_PATH = './models/faster_rcnn_models/'\
#     'vgg16_faster_rcnn_iter_490000.ckpt'
# WEIGHT_PATH = './models/vgg16.weights'
MODEL_PATH = './models/vgg16_faster_rcnn_iter_490000.ckpt'
WEIGHT_PATH = './models/vgg16.weights'
INPUT = './images'

def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(
        description='Index image to elasticsearch')
    parser.add_argument('--weight', dest='weight',
                        help='weight to test',
                        default=WEIGHT_PATH, type=str)
    parser.add_argument('--model_path', dest='model_path',
                        help='path to the model',
                        default=MODEL_PATH, type=str)

    parser.add_argument('--input', dest='input',
                        help='Input image folder',
                        default=INPUT, type=str)

    parser.add_argument('--es_host', dest='es_host',
                        help='es sever host',
                        default='localhost', type=str)
    parser.add_argument('--es_index', dest='es_index',
                        help='index name',
                        default='images', type=str)
    parser.add_argument('--es_type', dest='es_type',
                        help='index type',
                        default='obj', type=str)
    parser.add_argument('--es_port', dest='es_port',
                        help='es server port',
                        default=9200, type=int)

    args = parser.parse_args()

    if not args.input:
        parser.error('Input folder not given')
    return args


def create_doc(im_src, tag):
    doc = {}
    doc['im_src'] = im_name
    doc['im_tag'] = tag
    return doc


if __name__ == '__main__':
    args = parse_args()

    f_val = open("synset_words.txt", "r")
    valDict = {}
    for i, line in enumerate(f_val.readlines()):
        words = (line.replace(',', '')).split()
        valDict[i] = words[1:]
    f_val.close()

    f_imgs = open("val.txt", "r")
    imgDict = {}
    for line in f_imgs.readlines():
        words = (line.replace('\n', '')).split()
        value = int(words[1])
        imgDict[words[0]] = valDict[value]
    f_imgs.close()

    # create elasticsearch client
    es = Elasticsearch(hosts='{}:{}'.format(args.es_host, args.es_port))

    # load images
    images = glob.glob(args.input + "/*")
    print "image number:", len(images)

    bulk, actions = [], []
    count, num_docs = 0, 0
    es_index = args.es_index
    es_type = args.es_type

    num_imgs = len(images)
    for im_path in images:
        count += 1
        if count % 100 == 0:
            logger.info('Processing image {}/{}'.format(count, num_imgs))

        # index document
        im_name = im_path.split('/')[-1]
        doc = {}
        doc = create_doc(im_name, imgDict[im_name])
        num_docs += 1

        # create index action
        action = {
            "_index": es_index,
            "_type": es_type,
            "_source": doc
        }
        actions.append(action)

        if len(actions) == 1000:
            logger.info('Bulking {} docs to sever, indexed {}'
                        .format(len(actions), num_docs))
            helpers.bulk(es, actions)
            del actions[:]

    if len(actions) > 0:
        helpers.bulk(es, actions)
        logger.info('Bulking {} docs to sever,  total {}'
                    .format(len(actions), num_docs))


